import pandas as pd
from data import Dataset
from utils import *
from time import time,sleep
import os
import re
import warnings
from multiprocessing import Pool as ThreadPool 
from multiprocessing import Lock

class Speculate():
    def __init__(self):
        self.df = pd.DataFrame() # 预测数据
        self.data = Dataset()
        
    
    def tidy_5min_data(self,src,dst):
        # 按年整理通信达下载的5分钟数据
        assert os.path.exists(src),'找不到原始文件夹'
        if not os.path.exists:
            os.mkdir(dst)
        # 待遍历的文件夹队列
        s_dir = [src,]
        while(len(s_dir)>0):
            last_dir = s_dir.pop()
            list_files = os.listdir(last_dir)
            # 显示进度
            file_num = len(list_files)
            floder_num = len(s_dir)
            fin = 0
            # 遍历当前文件夹文件
            for file in list_files:
                file_path = os.path.join(last_dir,file)
                # 如果是文件夹则入队列
                if os.path.isdir(file_path):
                    s_dir.append(file_path)
                    continue
                # 如果是文件，看下文件名是否符合通信达命名
                pattern = r'([A-Z]{2})#(\d{6}).([a-z]{3})'
                if not re.match(pattern,file):
                    continue
                # 读取数据
                # 通信达数据格式，YYYY-MM-dd,hhmm,开,高,低,收,量,额，无列名，gbk编码,末尾有中文标注要去除
                # 新增列名，trade_date,trade_min,open,high,low,close,vol,amount
                col_names = ['trade_date','trade_min','open','high','low','close','vol','amount']
                df = pd.read_csv(file_path,encoding='gbk',names = col_names,header=None).dropna()
                # 排号序，把年份取出来，按不同年份保存
                df = df.sort_values(by='trade_date',ascending=True).reset_index(drop=True)
                years = df['trade_date'].apply((lambda x : x.split('-')[0]))
                for year in years.unique():
                    year_data = df[years==year]
                    # 文件名及保存路径
                    name = ''.join(file.split('.')[:-1])
                    save_path = os.path.join(dst,f'{year}\{name}.csv')
                    if os.path.exists(save_path):
                        # 之前有保存的，读取合并去重再保存
                        last_data = pd.read_csv(save_path)
                        year_data = pd.concat([last_data,year_data])
                        year_data.drop_duplicates(subset=['trade_date','trade_min'],keep='first',inplace=True)
                        year_data.sort_values(by=['trade_date','trade_min'],ascending=True).reset_index(drop=True)
                        year_data.to_csv(save_path)
                    else:
                        if not os.path.exists(os.path.join(dst,year)):
                            os.mkdir(os.path.join(dst,year))
                        year_data.to_csv(save_path)
                # 显示进度
                fin +=1
                print(f'\r正在处理文件夹: {last_dir}, 进度: {fin}/{file_num}, 剩余{floder_num}个文件夹处理',end='',flush=True)
    
    def mk_data(self,src,dst,features_days=1,features_fun=None):
        """
        自动生成训练数据
        features_days: 以前features_days天的五分钟数据作为训练特征
        features_fun: 特征处理函数,根据不同的特征组织需求实现不同的特征函数,如果该参数为None,则保存直接数据
        src: 源数据地址
        dst: 处理后目标地址
        """
        assert os.path.exists(src),'找不到原始文件夹'
        assert os.path.exists(dst),'目标文件夹不存在'
        # 下面多多进程处理初始化
        warnings.simplefilter(action='ignore', category=FutureWarning)
        self.name_index = 0 # 用于输出文件名的，每10000行做一个文件
        self.file_lock = Lock() # 锁
        propool = ThreadPool(4) # 进程池
        # 待遍历的文件夹队列
        s_dir = [src,]
        while(len(s_dir)>0):
            last_dir = s_dir.pop()
            list_files = os.listdir(last_dir)
            # 显示进度
            file_num = len(list_files)
            floder_num = len(s_dir)
            fin = 0
            # 遍历当前文件夹文件
            for file in list_files:
                file_path = os.path.join(last_dir,file)
                # 如果是文件夹则入队列
                if os.path.isdir(file_path):
                    s_dir.append(file_path)
                    continue
                # 如果是文件，看下文件名是否符合通信达命名SZ#123456.csv
                pattern = r'([A-Z]{2})#(\d{6}).([a-z]{3})'
                if not re.match(pattern,file):
                    continue
                # 读取数据，按照日期分组,然后把转化为list
                df = pd.read_csv(file_path,index_col=0).sort_values(by=['trade_date','trade_min'],ascending=True).reset_index(drop=True)
                group = df.groupby(['trade_date'])
                group = list(group)
                # 设置长度，每次取features_days天,后面不足features_days天的数据丢弃
                length = len(group) - len(group) % features_days
                for i in range(0,length,features_days):
                    df = pd.DataFrame()
                    # 数据按照features_days合并
                    for item in group[i:i+features_days]:
                        df = pd.concat([df,item[1]])
                    df = df.reset_index(drop=True)
                    if features_fun:
                        trade_code = file.split('.')[0].split('#')
                        trade_code = trade_code[1] + '.' + trade_code[0]
                        # propool.apply_async(features_fun,(trade_code,df,dst,))
                        features_fun(trade_code,df,dst)
                    else:
                        file_name = file.split('.')[0] + '_' + str(time()*100).split('.')[0][2:] + '.csv'
                        df.to_csv(os.path.join(dst,file_name))
                # 显示进度
                fin +=1
                
                if fin == 200:
                    propool.close()
                    propool.join()
                    return
                
                print(f'\r正在处理文件夹: {last_dir}, 进度: {fin}/{file_num}, 剩余{floder_num}个文件夹处理',end='',flush=True)
        propool.close()
        propool.join()    
    
    def fun_close_vol(self,ts_code,df,dst):
        """
        特征训练数据生成函数
        特征: 分时数据的收盘数据和成交量作为特征
        标签: 后一个交易日的涨跌
        """
        # 获取最后的交易日，然后查询下一个交易日涨跌作为标签
        df_date = pd.unique(df['trade_date'])
        df_date.sort()
        df_date = ''.join(df_date[-1].split('-'))
        next_date = Trade_date(df_date).next_trade_day()
        target,target_date = next(next_date),next(next_date) # 当天是交易日，所以第一次返回的是当天
        # 如果这一天为None，说明交易日在未来，则生成预测数据
        # 查询目标交易日的涨跌作为标签
        # 手动设置涨跌幅阈值，pct_chg为涨跌幅百分点数
        while(True):
            target = self.data.get_by_code_date(ts_code,target_date)
            if len(target)>0:
                target = target.iloc[0]['pct_chg']
                target = 1 if target >= 1 else 0
                break
            target_date = next(next_date)
            if not target_date:
                break
        # 组织整理数据
        df = df.sort_values(by=['trade_date','trade_min'],ascending=True) # 先排个序
        save_df = {} # 数据
        col_name = [] # 列名
        i = 0 # 区分列名
        for item in df.itertuples():
            save_df['close'+str(i)] = [item.close]
            save_df['vol'+str(i)] = [item.vol]
            col_name.append('close'+str(i))
            col_name.append('vol'+str(i))
            i += 1
        # 添加标签
        col_name.append('target')
        save_df['target']=str(target)
        save_df = pd.DataFrame(save_df,columns=col_name)
        # 等待锁
        #self.file_lock.acquire()# 等待锁
        # 设置文件名然后保存
        save_name = '5min_train_{}.csv'.format(str(int(self.name_index/10000)))
        self.name_index += 1
        save_path = os.path.join(dst,save_name)
        if not os.path.exists(save_path):
            save_df.to_csv(save_path,index=False)
        else:
            save_df.to_csv(save_path,mode='a',index=False,header=False)
        # self.file_lock.release() # 解锁
        
             
        
                    
                
                
        
        